"""
@description: Utility functions and dataclasses for configuration and logging.
@author: yangkw
@date: 2025-05-02
"""
# -*- coding: utf-8 -*-
import yaml
import logging
from pathlib import Path
from dataclasses import dataclass, field, fields
from typing import Optional, List, Union

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 使用 dataclasses 定义训练参数结构，便于类型检查和IDE提示
# 注意：这里的字段需要与 training_args.yaml 中的键名保持一致
@dataclass
class ModelArguments:
    """
    与模型和分词器相关的参数
    """
    model_name_or_path: str = field(
        metadata={"help": "预训练模型名称或路径 (Hugging Face Hub 或本地)。"}
    )
    tokenizer_name_or_path: Optional[str] = field(
        default=None, metadata={"help": "预训练分词器名称或路径。如果为 None，则使用 model_name_or_path。"}
    )
    # 可以根据需要添加更多模型相关参数，例如 trust_remote_code=True

@dataclass
class DataTrainingArguments:
    """
    与数据处理和训练相关的参数
    """
    train_file: Optional[str] = field(
        default=None, metadata={"help": "训练数据文件路径。"}
    )
    validation_file: Optional[str] = field(
        default=None, metadata={"help": "验证数据文件路径。"}
    )
    text_column: str = field(
        default="text", metadata={"help": "数据文件中包含文本的列名。"}
    )
    max_seq_length: int = field(
        default=512, metadata={"help": "分词后输入序列的最大长度。"}
    )
    # 可以添加数据预处理相关的其他参数，例如是否移除长序列等

@dataclass
class QuantizationArguments:
    """
    与模型量化相关的参数 (BitsandBytes)
    """
    load_in_4bit: bool = field(
        default=False, metadata={"help": "是否以 4-bit 加载模型。"}
    )
    bnb_4bit_quant_type: str = field(
        default="nf4", metadata={"help": "4-bit 量化类型 ('nf4' 或 'fp4')。"}
    )
    bnb_4bit_compute_dtype: str = field(
        default="bfloat16", metadata={"help": "4-bit 计算数据类型 ('float16' 或 'bfloat16')。"}
    )
    bnb_4bit_use_double_quant: bool = field(
        default=False, metadata={"help": "是否对 4-bit 量化使用双重量化。"}
    )

@dataclass
class PeftArguments:
    """
    与 PEFT (如 LoRA) 相关的参数
    """
    use_peft: bool = field(
        default=False, metadata={"help": "是否使用 PEFT (如 LoRA) 进行微调。"}
    )
    lora_r: int = field(
        default=8, metadata={"help": "LoRA 秩 (rank)。"}
    )
    lora_alpha: int = field(
        default=16, metadata={"help": "LoRA alpha 参数。"}
    )
    lora_dropout: float = field(
        default=0.05, metadata={"help": "LoRA 层的 dropout 概率。"}
    )
    lora_target_modules: Optional[Union[List[str], str]] = field(
        default=None, metadata={"help": "要应用 LoRA 的模块名称列表或单个模块名称。"}
    )
    # 可以添加其他 PEFT 参数，例如 bias 类型等

@dataclass
class TrainingArguments:
    """
    Hugging Face Trainer 的核心训练参数 (部分常用参数)
    注意：实际训练脚本中我们会直接使用 transformers.TrainingArguments，
    这里定义是为了方便从 YAML 加载并传递给 Trainer。
    """
    output_dir: str = field(
        metadata={"help": "训练输出（检查点、日志等）的目录。"}
    )
    num_train_epochs: float = field(
        default=3.0, metadata={"help": "总的训练轮数。"}
    )
    per_device_train_batch_size: int = field(
        default=8, metadata={"help": "每个 GPU/TPU 核心的训练批量大小。"}
    )
    per_device_eval_batch_size: int = field(
        default=8, metadata={"help": "每个 GPU/TPU 核心的评估批量大小。"}
    )
    gradient_accumulation_steps: int = field(
        default=1, metadata={"help": "在执行反向传播/更新之前累积梯度的更新步数。"}
    )
    learning_rate: float = field(
        default=5e-5, metadata={"help": "初始学习率 (AdamW 优化器)。"}
    )
    lr_scheduler_type: str = field(
        default="linear", metadata={"help": "使用的学习率调度器类型。"}
    )
    warmup_ratio: float = field(
        default=0.0, metadata={"help": "用于线性预热的步数占总训练步数的比例。"}
    )
    weight_decay: float = field(
        default=0.0, metadata={"help": "应用于除所有偏差和 LayerNorm 权重之外的所有层的权重衰减。"}
    )
    optim: str = field(
        default="adamw_torch", metadata={"help": "使用的优化器。"}
    )
    max_grad_norm: float = field(
        default=1.0, metadata={"help": "最大梯度范数 (用于梯度裁剪)。"}
    )
    logging_steps: int = field(
        default=500, metadata={"help": "每隔多少步记录一次日志。"}
    )
    save_strategy: str = field(
        default="steps", metadata={"help": "模型保存策略 ('steps' 或 'epoch')。"}
    )
    save_steps: int = field(
        default=500, metadata={"help": "如果 save_strategy='steps'，则每隔多少步保存一次检查点。"}
    )
    save_total_limit: Optional[int] = field(
        default=None, metadata={"help": "限制检查点的总数。删除旧的检查点。"}
    )
    evaluation_strategy: str = field(
        default="no", metadata={"help": "评估策略 ('no', 'steps', 'epoch')。"}
    )
    eval_steps: Optional[int] = field(
        default=None, metadata={"help": "如果 evaluation_strategy='steps'，则每隔多少步进行一次评估。"}
    )
    report_to: Optional[List[str]] = field(
        default=None, metadata={"help": "结果和日志报告的目标列表 (例如 'wandb', 'tensorboard')。"}
    )
    seed: int = field(
        default=42, metadata={"help": "用于复现性的随机种子。"}
    )
    bf16: bool = field(
        default=False, metadata={"help": "是否使用 bf16 混合精度训练 (需要 Ampere 或更新架构)。"}
    )
    fp16: bool = field(
        default=False, metadata={"help": "是否使用 fp16 混合精度训练。"}
    )
    gradient_checkpointing: bool = field(
        default=False, metadata={"help": "是否使用梯度检查点来节省内存，但会牺牲训练速度。"}
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "是否覆盖 datasets 库的缓存文件。"}
    )
    # DeepSpeed 配置由 Accelerate 自动处理，不需要在这里指定
    deepspeed_config: str = field(
        default=None, metadata={"help": "deepspeed_config文件路径"}
    )
    resume_from_checkpoint: Optional[str] = field(
        default=None, metadata={"help": "用于恢复训练的检查点路径。"}
    )
    # Early Stopping related arguments
    early_stopping_patience: int = field(
        default=0, metadata={"help": "早停的耐心值。如果大于0，则启用早停。"}
    )
    load_best_model_at_end: bool = field(
        default=False, metadata={"help": "训练结束后是否加载指标最佳的模型检查点。"}
    )
    metric_for_best_model: Optional[str] = field(
        default=None, metadata={"help": "用于确定最佳模型的指标名称 (例如 'eval_loss')。"}
    )
    greater_is_better: Optional[bool] = field(
        default=None, metadata={"help": "指定的指标是否越大越好。"}
    )

@dataclass
class WandbArguments:
    """
    与 Wandb 日志记录相关的参数
    """
    wandb_project: Optional[str] = field(
        default="llm-finetune-project", metadata={"help": "Wandb 项目名称。"}
    )
    wandb_entity: Optional[str] = field(
        default=None, metadata={"help": "Wandb 实体 (用户名或团队名)。"}
    )
    wandb_name: Optional[str] = field(
        default=None, metadata={"help": "Wandb 运行名称 (留空则自动生成)。"}
    )


def load_config(config_path: Union[str, Path]) -> dict:
    """
    从 YAML 文件加载配置。

    Args:
        config_path: YAML 配置文件的路径。

    Returns:
        包含配置项的字典。
    """
    config_path = Path(config_path)
    if not config_path.is_file():
        logger.error(f"配置文件未找到: {config_path}")
        raise FileNotFoundError(f"配置文件未找到: {config_path}")

    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        logger.info(f"成功加载配置文件: {config_path}")
        return config
    except Exception as e:
        logger.error(f"加载配置文件失败: {config_path}, 错误: {e}")
        raise

def parse_args_from_yaml(config_path: Union[str, Path]):
    """
    从 YAML 文件加载配置，并将其解析到定义的 dataclass 结构中。

    Args:
        config_path: YAML 配置文件的路径。

    Returns:
        包含所有解析后参数对象的元组 (model_args, data_args, training_args, quant_args, peft_args, wandb_args)。
    """
    config = load_config(config_path)

    arg_classes = {
        "model": ModelArguments,
        "data": DataTrainingArguments,
        "training": TrainingArguments, # 注意：这里的 training 对应 YAML 中的大部分训练参数
        "quantization": QuantizationArguments,
        "peft": PeftArguments,
        "wandb": WandbArguments
    }

    parsed_args = {}

    # 遍历所有定义的 dataclass
    all_defined_fields = set()
    for arg_class in arg_classes.values():
        for f in fields(arg_class):
            all_defined_fields.add(f.name)

    # 将 YAML 中的顶级键值对分配给对应的 dataclass
    init_args = {arg_name: {} for arg_name in arg_classes}
    consumed_keys = set()

    for key, value in config.items():
        assigned = False
        # 尝试将键分配给对应的 dataclass
        for arg_name, arg_class in arg_classes.items():
            class_fields = {f.name for f in fields(arg_class)}
            if key in class_fields:
                init_args[arg_name][key] = value
                consumed_keys.add(key)
                assigned = True
                break # 一个键只属于一个 dataclass

        # 特殊处理 training 参数，因为 YAML 中很多顶级键属于 TrainingArguments
        if not assigned and key in {f.name for f in fields(TrainingArguments)}:
             init_args["training"][key] = value
             consumed_keys.add(key)
             assigned = True

        # if not assigned:
        #     logger.warning(f"配置文件中的键 '{key}' 未在任何参数类中定义，将被忽略。")


    # 检查是否有未在 YAML 中定义的必需字段 (没有默认值的字段)
    # for arg_name, arg_class in arg_classes.items():
    #     for f in fields(arg_class):
    #         if f.default is field.MISSING and f.default_factory is field.MISSING and f.name not in init_args[arg_name]:
    #             raise ValueError(f"配置文件 '{config_path}' 中缺少必需参数 '{f.name}' (属于 {arg_class.__name__})。")

    # 实例化 dataclass 对象
    try:
        parsed_args["model_args"] = ModelArguments(**init_args["model"])
        parsed_args["data_args"] = DataTrainingArguments(**init_args["data"])
        # 对于 TrainingArguments，需要合并 YAML 顶级键和 'training' 部分 (如果存在)
        training_combined_args = init_args.get("training", {})
        # training_combined_args.update(init_args.get("training_specific", {})) # 如果 YAML 里有 'training' section
        parsed_args["training_args"] = TrainingArguments(**training_combined_args)
        parsed_args["quant_args"] = QuantizationArguments(**init_args["quantization"])
        parsed_args["peft_args"] = PeftArguments(**init_args["peft"])
        parsed_args["wandb_args"] = WandbArguments(**init_args["wandb"])
    except TypeError as e:
         logger.error(f"解析配置文件时发生类型错误，请检查参数名称和类型是否匹配: {e}")
         raise

    # 检查 YAML 中是否有未被使用的键
    unused_keys = set(config.keys()) - consumed_keys
    if unused_keys:
        logger.warning(f"配置文件中的以下键未被使用: {', '.join(unused_keys)}")


    return (
        parsed_args["model_args"],
        parsed_args["data_args"],
        parsed_args["training_args"],
        parsed_args["quant_args"],
        parsed_args["peft_args"],
        parsed_args["wandb_args"]
    )

# 示例用法 (通常在主脚本中调用)
if __name__ == '__main__':
    # 假设 config 文件在父目录的 config/ 下
    script_dir = Path(__file__).parent
    project_root = script_dir.parent
    example_config_path = project_root / "config" / "training_args.yaml"

    if example_config_path.exists():
        try:
            model_args, data_args, training_args, quant_args, peft_args, wandb_args = parse_args_from_yaml(example_config_path)
            logger.info("配置文件解析成功!")
            logger.info(f"Model Arguments: {model_args}")
            logger.info(f"Data Arguments: {data_args}")
            logger.info(f"Training Arguments: {training_args}")
            logger.info(f"Quantization Arguments: {quant_args}")
            logger.info(f"PEFT Arguments: {peft_args}")
            logger.info(f"Wandb Arguments: {wandb_args}")

            # 访问具体参数示例
            print(f"\n模型名称: {model_args.model_name_or_path}")
            print(f"训练文件: {data_args.train_file}")
            print(f"输出目录: {training_args.output_dir}")
            print(f"是否使用4bit量化: {quant_args.load_in_4bit}")
            print(f"是否使用PEFT: {peft_args.use_peft}")
            print(f"Wandb项目: {wandb_args.wandb_project}")

        except (FileNotFoundError, ValueError, TypeError) as e:
            logger.error(f"处理配置文件时出错: {e}")
    else:
        logger.warning(f"示例配置文件未找到: {example_config_path}，无法执行示例解析。")
